import os
import pandas as pd
import json
import openai
import openpyxl
from openpyxl.utils import get_column_letter
from openpyxl.styles import Alignment
from typing import List

###############################################################################
# 生成excel表格具体内容
def openai_setup():
    ############################################################################## 需改
    # 确保设置了环境变量 OPENAI_API_KEY
    # 设置OpenAI API密钥
    openai.api_key = '你的apikey' # 替换为你的OpenAI API密钥
    openai.api_base = "https://www.dmxapi.com/v1"  # 可以替换为其他API端点
    ##############################################################################

def generate_remarks(
    excel_path: str,
    output_path: str,
    country: str,
    subject: str,
    code: str,
    model: str = "gpt-4o"
):
    """
    读取 Excel 文件，从第三行开始（索引 2）根据第一列的值，与指定的国家和描述对象，
    使用大模型生成备注内容并写入第五列。若无法检索到相关信息，则填写 'NULL'.

    参数：
    excel_path: 输入 Excel 文件路径
    output_path: 输出结果 Excel 文件路径
    country: 国家名称
    subject: 描述对象名称
    model: 要调用的模型名称，默认为 gpt-4o
    """

    openai_setup()  # 设置 OpenAI API 密钥和基础 URL

    # 读取 Excel
    df = pd.read_excel(excel_path)

    if df.iat[1,0] == '名称':
        df.iat[1, 4] = subject  # 第二行的第五列

    total_rows = len(df) - 2  # 从第三行开始
    print(f"需要填写备注的行数: {total_rows}")

    # 遍历从第三行开始的每一行
    for idx in range(2, len(df)):
        item = df.iat[idx, 0]  # 第一列的值
        item2 = df.iat[idx, 4]  # 第五列的值

        staticListCODE = ['分类码', '图片']
        staticListNULL = ['区域代码', '国家或行业标准编码' ,'军队内部编码', '军用基础地理信息编码', '地理标示', '关联地图图种及比例尺', '外挂表名称', '视频', '地图', '关联地图图种及比例尺', '视频数据']
        if item in staticListCODE:
            df.iat[idx, 4] = code  # 第五列
            # print(f"行 {idx+1} - 项目: {item} -> 备注: {code}")
            continue
        elif item == '顺序码':
            df.iat[idx, 4] = '信息采集自动生成'  # 第五列
            # print(f"行 {idx+1} - 项目: {item} -> 备注: 信息采集自动生成")
            continue
        elif item in staticListNULL:
            df.iat[idx, 4] = 'NULL'  # 第五列
            # print(f"行 {idx+1} - 项目: {item} -> 备注: NULL")
            continue


        prompt = (
            f"请根据以下信息生成一段简洁明了的备注，用于补充说明或背景介绍：\n"
            f"国家：{country}\n"
            f"{subject}：{item}\n"
            f"{item2}\n\n"
            f"请优先基于公开资料或常识生成真实可信的数据和内容；\n"
            f"若确实检索不到可考来源，可合理推测\n"
            f"如果实在无法推理同时不是相干信息，请仅回复 NULL。\n"
)
        try:
            resp = openai.ChatCompletion.create(
                model=model,
                messages=[
                    {"role": "system", "content": "你是一个信息检索与补充的助手，能结合事实与合理推测并给出数据和资料。描述尽量简洁，每次生成不超过150字。"},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.0,
                max_tokens=500
            )
            remark = resp.choices[0].message.content.strip()
        except Exception as e:
            print(f"行 {idx+1} 调用 API 时出错: {e}")
            remark = "NULL"

        # 如果模型返回为空或包含无法识别内容，也填写 NULL
        if not remark or remark.upper() == "N/A":
            remark = "NULL"

        df.iat[idx, 4] = remark  # 第五列
        # print(f"行 {idx+1} - 项目: {item} -> 备注: {remark}")

    # 确保输出文件的父目录存在
    output_dir = os.path.dirname(output_path)
    os.makedirs(output_dir, exist_ok=True)

    # 保存到新的 Excel 文件
    df.to_excel(output_path, index=False)
    print(f"备注已写入并保存到 {output_path}")


###############################################################################
# 下面的函数用于使用 GPT-4o 模型进行广泛搜索，返回关于指定国家和描述对象的列表
def search_by_gpt4o(
    country: str,
    description: str,
    count: int = 5,
    model: str = "gpt-4o",
    max_retries: int = 2,
    temperature: float = 0.3
) -> List[str]:
    """
    使用 GPT-4o “模拟”广泛搜索，返回关于指定国家和描述对象的列表。
    自动剥离返回文本中的 Markdown 代码块（```json ... ```），确保能被 json.loads 解析。
    """

    openai_setup()  # 设置 OpenAI API 密钥和基础 URL

    system_msg = (
        "你是一个强大的“虚拟搜索引擎”。"
        "请在指定的国家范围内，对给定的描述对象进行广泛搜索，"
       f"如果有对应内容，请返回相关的对象名称（或者自命名），保证至少找到 {count} 个，可以以每年为一个单位，生成时带上年份，年份不能超过2025年，"
        "如果没有找到相关内容，请返回空列表。"
        "并以纯 JSON 数组的格式返回找到的对象名称列表，"
        "不要包含任何多余的解释或文本。"
    )

    print(f'国家：{country}')
    print(f'要搜索的对象描述：{description}')


    user_msg = (
        f"国家：{country}\n"
        f"要搜索的对象描述：{description}\n\n"
        f"请返回至少 {count} 条结果，用 JSON 列表的形式，例如：\n" 
        '["对象1", "对象2", ...]'
        '对象的名称请使用中文。'
    )

    for attempt in range(1, max_retries + 1):
        resp = openai.ChatCompletion.create(
            model=model,
            messages=[
                {"role": "system",  "content": system_msg},
                {"role": "user",    "content": user_msg},
            ],
            temperature=temperature,
            timeout = 100,
        )
        text = resp.choices[0].message.content.strip()

        # —— 新增：去除 Markdown 代码块 包裹
        if text.startswith("```"):
            lines = text.splitlines()
            # 去掉首尾 ``` 或 ```json
            if lines[0].startswith("```"):
                lines = lines[1:]
            if lines and lines[-1].startswith("```"):
                lines = lines[:-1]
            text = "\n".join(lines).strip()
        # ————————————

        try:
            results = json.loads(text)
            if isinstance(results, list) and all(isinstance(x, str) for x in results):
                return results
            else:
                raise ValueError("返回格式不是字符串列表")
        except Exception as e:
            if attempt == max_retries:
                raise RuntimeError(f"解析模型输出为列表失败：{e}\n清理后文本：\n{text}")
            # 重试提示
            user_msg = (
                "上一次返回的不是合法的 JSON 字符串数组，请仅返回 JSON 列表，"
                "不要多余说明。\n\n"
                "国家：" + country + "\n"
                "要搜索的对象描述：" + description
            )

    return []

###############################################################################
# 调整excel格式
def modify_excel(output_path):
    # 加载 Excel 文件
    wb = openpyxl.load_workbook(output_path)
    ws = wb.active  # 默认处理第一个工作表

    # 清除第一行第2到第5列的内容（B1 到 E1）
    for col in range(2, 6):  # B=2, E=5
        ws.cell(row=1, column=col).value = None

    # 合并第一行第1到第5列（A1:E1）
    ws.merge_cells('A1:E1')

    # 可选：设置合并单元格的对齐方式为居中
    merged_cell = ws['A1']
    merged_cell.alignment = Alignment(horizontal='center', vertical='center')

    # 保存修改
    wb.save(output_path)
    print(f"已成功处理文件：{output_path}")

###############################################################################
# 获取文件名和格式文件
def extract_code_object_from_folder(folder_path, start_idx, end_idx):
    """
    从文件夹中指定索引区间的 Excel 文件名中提取 code 和 object。

    参数:
        folder_path (str): 文件夹路径
        start_idx (int): 起始索引（包含）
        end_idx (int): 结束索引（不包含）

    返回:
        List[Tuple[str, str, str]]: 每项是 (文件名, code, object)
    """
    results = []
    # 获取所有 .xlsx 文件并排序
    all_files = sorted([f for f in os.listdir(folder_path) if f.endswith('.xlsx')])

    # 截取指定区间的文件名
    for fname in all_files[start_idx:end_idx]:
        try:
            code, obj = fname.split('-')[0], fname.split('-')[1].rsplit('.', 1)[0]
            results.append((fname, code, obj))
        except Exception as e:
            print(f"跳过无效文件名格式: {fname}，错误：{e}")
    return results


# if __name__ == "__main__":

#     country = "加勒比海地区"
#     folder_path = "output"
#     source_path = r"C:\Users\xsr\Desktop\赵的数据"


def solution(country, folder_path, source_path):

    # 指定开始和结束
    start = 0
    end = 1000
    i = start

    ################################################################################ 需改
    # 指定生成的个数
    digitarr = [ 150,  151, 158 , 150, 152 ]
    # 指定生成的excel模板对象名称
    categories = [
    "军事同盟",
    "军事冲突",
    "投资外汇",
    "对外贸易",
    "经济共同体"
    ]
    ################################################################################
    

    info_list = extract_code_object_from_folder(source_path, start, end)
    for fname, code, object in info_list:
        excel_path = source_path + '/' + fname
        print(f"文件: {excel_path} -> code: {code}, object: {object}")

        try:
            # output_object_result = search_by_gpt4o(country, object, digitarr[i])
            output_object_result = search_by_gpt4o(country, categories[i], digitarr[i])
            print("搜索对象应有个数：", digitarr[i])
            print("搜索结果：", output_object_result)
        except Exception as err:
            print("出错：", err)
        
        
        i += 1
        count = len(output_object_result)
        print("搜索结果数量：", count)

        for j in range(count):
            output_path = folder_path + "/" + code + "/" + f"{code}{j+1:03d}.xlsx"
            print("生成的文件路径：", output_path)
            generate_remarks(
                excel_path=excel_path,
                output_path=output_path,
                country=country,
                subject=output_object_result[j],
                code=f"{code}{j+1:03d}",
                model="gpt-4o"
            )
            modify_excel(output_path=output_path)
            

if __name__ == "__main__":
    # country = "加勒比海地区"
    # folder_path = "output"

    ############################################################################## 需改
    # 指定国家和文件夹路径
    source_path = r"F:\code\automatic-data-generation\data_generate\基底数据表"

    countrys = ['土库曼斯坦地区', '加勒比海地区']
    folder_paths = [ 'output/土库曼斯坦关系数据', 'output/加勒比海地区关系数据'] # 输出路径
    ##############################################################################


    for i in range(len(countrys)):
        country = countrys[i]
        folder_path = folder_paths[i]
        print("国家：", country)
        print("文件夹：", folder_path)
        solution(country, folder_path, source_path)

